/*
 * Copyright (c) 2022. China Mobile (SuZhou) Software Technology Co.,Ltd. All rights reserved.
 * Lakehouse is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package com.chinamobile.cmss.lakehouse.core.queue.resource;

import com.chinamobile.cmss.lakehouse.common.dto.ResourceCheckResult;
import com.chinamobile.cmss.lakehouse.common.dto.yunikorn.ConfigQueue;
import com.chinamobile.cmss.lakehouse.common.dto.yunikorn.ConfigResources;
import com.chinamobile.cmss.lakehouse.common.dto.yunikorn.ConfigResources.ConfigResource;
import com.chinamobile.cmss.lakehouse.common.dto.yunikorn.QueuesCapacity;
import com.chinamobile.cmss.lakehouse.common.enums.ResourceSizeEnum;
import com.chinamobile.cmss.lakehouse.common.kubernetes.ComputeResource;
import com.chinamobile.cmss.lakehouse.common.kubernetes.ResourceSizeConfig;
import com.chinamobile.cmss.lakehouse.core.queue.QueueResourceService;
import com.chinamobile.cmss.lakehouse.core.queue.ResourceManager;

import java.util.HashMap;
import java.util.Map;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
@Slf4j
public class QueueResourceServiceImpl implements QueueResourceService {

    @Autowired
    private Map<String, ResourceManager> manager;

    @Override
    public ResourceCheckResult runnable(String queue, int cpu, int memory) {
        try {
            boolean check = manager.get(ResourceSizeConfig.resourceManager).runnable(queue, cpu, memory);
            return new ResourceCheckResult(null, check);
        } catch (Exception ex) {
            log.error("Cannot check runnable: {}", ex.getMessage());
            return new ResourceCheckResult(ex.getMessage(), false);
        }
    }

    @Override
    public synchronized boolean createYunikornQueue(String queueName,
                                                    ResourceSizeEnum resourceSize) {
        try {
            ComputeResource computeResource = ResourceSizeConfig.getResource(resourceSize);
            String strategy = ResourceSizeConfig.resourceStrategy;
            ConfigQueue configQueue = null;
            if ("fifo".equals(strategy) || "fair".equals(strategy) || "stateaware".equals(strategy)) {
                configQueue = generateYunikornQueue(queueName, computeResource, strategy);
            } else {
                configQueue = generateYunikornQueue(queueName, computeResource, "fifo");
            }
            manager.get(ResourceSizeConfig.resourceManager).createQueue(configQueue);
        } catch (Throwable ex) {
            log.error("Cannot create yunikorn queue, reason: {}", ex.getMessage());
            return false;
        }
        return true;
    }

    @Override
    public synchronized ResourceCheckResult resizeYunikornQueue(String queueName,
                                                                ResourceSizeEnum resourceSize) {
        try {
            ComputeResource computeResource = ResourceSizeConfig.getResource(resourceSize);
            int cpu = computeResource.getCpu() * 1000;
            int memory = computeResource.getMemory() * 1024;
            ConfigResources resources = new ConfigResources();
            ConfigResource guaranteed = new ConfigResource();
            guaranteed.setMemory(memory + "");
            guaranteed.setVcore(cpu + "");
            resources.setGuaranteed(guaranteed);
            resources.setMax(guaranteed);
            manager.get(ResourceSizeConfig.resourceManager).editConfigQueue(queueName, resources);
            return new ResourceCheckResult(null, true);
        } catch (Exception ex) {
            log.error("Cannot resize yunikorn queue: {}, reason: {}", queueName, ex.getMessage());
            return new ResourceCheckResult(ex.getMessage(), false);
        }

    }

    @Override
    public synchronized boolean deleteYunikornQueue(String queueName) {
        try {
            manager.get(ResourceSizeConfig.resourceManager).deleteQueue(queueName);
        } catch (Exception ex) {
            log.error("Cannot delete yunikorn queue, reason: {}", ex.getMessage());
            return false;
        }
        return true;
    }

    @Override
    public ResourceCheckResult checkResizeYunikornQueue(String queueName,
                                                        ResourceSizeEnum resourceSize) {
        ComputeResource computeResource = ResourceSizeConfig.getResource(resourceSize);
        int cpu = computeResource.getCpu() * 1000;
        int memory = computeResource.getMemory() * 1024;
        ConfigResources resources = new ConfigResources();
        ConfigResource guaranteed = new ConfigResource();
        guaranteed.setMemory(memory + "");
        guaranteed.setVcore(cpu + "");
        resources.setGuaranteed(guaranteed);
        resources.setMax(guaranteed);
        return manager.get(ResourceSizeConfig.resourceManager).checkEditQueue(queueName, resources);
    }

    @Override
    public ResourceCheckResult checkCreateYunikornQueue(String queueName,
                                                        ResourceSizeEnum resourceSize) {
        ComputeResource computeResource = ResourceSizeConfig.getResource(resourceSize);
        ConfigQueue configQueue = generateYunikornQueue(queueName, computeResource, "fifo");
        return manager.get(ResourceSizeConfig.resourceManager).checkCreateQueue(configQueue);
    }

    @Override
    public boolean checkExists(String queueName) {
        return manager.get(ResourceSizeConfig.resourceManager).existsQueueName(queueName);
    }

    @Override
    public QueuesCapacity getQueueCapacity() {
        return manager.get(ResourceSizeConfig.resourceManager).getQueuesCapacity();
    }

    private ConfigQueue generateYunikornQueue(String name, ComputeResource computeResource,
                                              String strategy) {
        // add queue
        int cpu = computeResource.getCpu() * 1000;
        int memory = computeResource.getMemory() * 1024;
        ConfigQueue addQueue = new ConfigQueue();
        addQueue.setName(name);
        ConfigResources resources = new ConfigResources();
        ConfigResource guaranteed = new ConfigResource();
        guaranteed.setMemory(memory + "");
        guaranteed.setVcore(cpu + "");
        resources.setGuaranteed(guaranteed);
        resources.setMax(guaranteed);
        addQueue.setResources(resources);
        Map<String, String> properties = new HashMap<>();
        properties.put("application.sort.policy", strategy);
        addQueue.setProperties(properties);
        return addQueue;
    }

}
